//图片添加水印

#include <stdio.h>
#include "libavformat/avformat.h"
#include "libavfilter/avfilter.h"
#include "libavfilter/buffersrc.h"
#include "libavfilter/buffersink.h"
#include "libavcodec/avcodec.h"

AVFrame *av_frame_malloc(int width, int height, int format) {
    AVFrame *frame = av_frame_alloc();
    frame->width = width;
    frame->height = height;
    frame->format = format;
    av_frame_get_buffer(frame, 1);

    return frame;
}

int save_frame_to_i420p(const char *filename, const AVFrame *frame) {
    if (frame->format != AV_PIX_FMT_YUV420P) {
        printf("The fromat isn't AV_PIX_FMT_YUV420P.\n");
        return -1;
    }

    FILE *outfile = fopen(filename, "wb");
    if (outfile == NULL) {
        printf("fopen %s failed.\n", filename);
        return -1;
    }

    size_t y_size = frame->width * frame->height;
    size_t u_size = y_size / 4;

    fwrite((const char*)frame->data[0], 1, y_size, outfile);
    fwrite((const char*)frame->data[1], 1, u_size, outfile);
    fwrite((const char*)frame->data[2], 1, u_size, outfile);

    fclose(outfile);

    return 0;
}

int save_frame_to_jpeg(const char *filename, const AVFrame *frame) {
    char error_msg_buf[256] = {0};
    AVCodec *jpeg_codec = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
    if (jpeg_codec == NULL) {
        return -1;
    }
    AVCodecContext *jpeg_codec_ctx = avcodec_alloc_context3(jpeg_codec);
    if (jpeg_codec_ctx == NULL) {
        return -2;
    }

    //这里不能使用YUV420P
    jpeg_codec_ctx->pix_fmt = AV_PIX_FMT_YUVJ420P;
    jpeg_codec_ctx->width = frame->width;
    jpeg_codec_ctx->height = frame->height;
    jpeg_codec_ctx->time_base = (AVRational){1, 25};
    jpeg_codec_ctx->framerate = (AVRational){25, 1};

    AVDictionary *encoder_opts = NULL;
    av_dict_set(&encoder_opts, "flags", "+qscale", 0);
    av_dict_set(&encoder_opts, "qmax", "2", 0);
    av_dict_set(&encoder_opts, "qmin", "2", 0);

    int ret = avcodec_open2(jpeg_codec_ctx, jpeg_codec, &encoder_opts);
    if (ret < 0) {
        printf("avcodec open failed:%s\n", av_make_error_string(error_msg_buf, sizeof(error_msg_buf), ret));
        avcodec_free_context(&jpeg_codec_ctx);
        return -3;
    }

    av_dict_free(&encoder_opts);

    AVPacket pkt;
    av_init_packet(&pkt);
    pkt.data = NULL;
    pkt.size = 0;

    ret = avcodec_send_frame(jpeg_codec_ctx, frame);
    if (ret < 0) {
        printf("Error: %s\n", av_make_error_string(error_msg_buf, sizeof(error_msg_buf), ret));
        avcodec_free_context(&jpeg_codec_ctx);
        return -4;
    }

    ret = 0;
    while (ret >= 0) {
        ret = avcodec_receive_packet(jpeg_codec_ctx, &pkt);
        if (ret == AVERROR(EAGAIN))
            continue;
        else if (ret == AVERROR_EOF) {
            ret = 0;
            break;
        }

        FILE *outfile = fopen(filename, "wb");
        if (outfile == NULL) {
            printf("fopen %s failed.\n", filename);
            ret = -1;
            break;
        }

        if (fwrite((char *)pkt.data, 1, pkt.size, outfile) == pkt.size)
            ret = 0;
        else {
            printf("fwrite %s failed.\n", filename);
            ret = -1;
        }

        fclose(outfile);

        ret = 0;
        break;
    }

    avcodec_free_context(&jpeg_codec_ctx);
    return ret;
}

AVFrame *get_frame_from_jpegfile(const char *filename) {
    int ret = 0;
    AVFrame *frame = NULL;
    AVFormatContext *format_ctx = NULL;
    if ((ret = avformat_open_input(&format_ctx, filename, NULL, NULL)) != 0) {
        char error_msg_buf[256] = {0};
        av_strerror(ret, error_msg_buf, sizeof(error_msg_buf));
        printf("Error: avformat_open_input failed: %s\n", error_msg_buf);
        return NULL;
    }

    avformat_find_stream_info(format_ctx, NULL);

    AVCodec *codec = NULL;
    int video_stream_idx = -1;
    video_stream_idx = av_find_best_stream(format_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, &codec, 0);
    if (video_stream_idx < 0) {
        goto cleanup_and_return;
    }

    AVCodecContext *av_codec_ctx = avcodec_alloc_context3(codec);
    ret = avcodec_open2(av_codec_ctx, codec, NULL);
    if (ret < 0) {
        goto cleanup_and_return;
    }

    AVPacket pkt;
    av_init_packet(&pkt);
    pkt.data = NULL;
    pkt.size = 0;
    ret = av_read_frame(format_ctx, &pkt);
    if (ret < 0) {
        goto cleanup_and_return;
    }

    ret = avcodec_send_packet(av_codec_ctx, &pkt);
    if (ret < 0) {
        goto cleanup_and_return;
    }

    frame = av_frame_alloc();
    ret = avcodec_receive_frame(av_codec_ctx, frame);
    if (ret < 0)
        av_frame_free(&frame);

cleanup_and_return:
    if (format_ctx != NULL)
        avformat_close_input(&format_ctx);

    if (av_codec_ctx != NULL)
        avcodec_free_context(&av_codec_ctx);

    return frame;
}

AVFilterContext *mainsrc_ctx = NULL;
AVFilterContext *logosrc_ctx = NULL;
AVFilterContext *resultsink_ctx = NULL;
AVFilterGraph *filter_graph = NULL;

int init_filters(const AVFrame* main_frame, const AVFrame* logo_frame, int x, int y) {
    int ret = 0;
    AVFilterInOut *inputs = NULL;
    AVFilterInOut *outputs = NULL;
    char filter_args[1024] = {0};

    filter_graph = avfilter_graph_alloc();
    if (filter_graph == NULL) {
        printf("Error: allocate filter graph failed.\n");
        return -1;
    }

    snprintf(filter_args, sizeof(filter_args),
             "buffer=video_size=%dx%d:pix_fmt=%d:time_base=1/25:pixel_aspect=%d/%d[main];"
             "buffer=video_size=%dx%d:pix_fmt=%d:time_base=1/25:pixel_aspect=%d/%d[logo];"
             "[main][logo]overlay=%d:%d[result];"
             "[result]buffersink",
             main_frame->width, main_frame->height, main_frame->format, main_frame->sample_aspect_ratio.num, main_frame->sample_aspect_ratio.den,
             logo_frame->width, logo_frame->height, logo_frame->format, logo_frame->sample_aspect_ratio.num, logo_frame->sample_aspect_ratio.den,
             x, y);

    ret = avfilter_graph_parse2(filter_graph, filter_args, &inputs, &outputs);
    if (ret < 0) {
        printf("Cannot parse graph.\n");
        return ret;
    }

    ret = avfilter_graph_config(filter_graph, NULL);
    if (ret < 0) {
        printf("Cannot configure graph.\n");
        return ret;
    }

    mainsrc_ctx = avfilter_graph_get_filter(filter_graph, "Parsed_buffer_0");
    logosrc_ctx = avfilter_graph_get_filter(filter_graph, "Parsed_buffer_1");
    resultsink_ctx = avfilter_graph_get_filter(filter_graph, "Parsed_buffersink_3");

    return 0;
}

int main_picture_mix_logo(AVFrame *main_frame, AVFrame *logo_frame, AVFrame *result_frame) {
    int ret = 0;
    ret = av_buffersrc_add_frame(mainsrc_ctx, main_frame);
    if (ret < 0)
        return ret;

    ret = av_buffersrc_add_frame(logosrc_ctx, logo_frame);
    if (ret < 0)
        return ret;

    ret = av_buffersink_get_frame(resultsink_ctx, result_frame);
    return ret;
}

int main(int argc, char *argv[])
{
    if (argc != 4) {
        printf("Usage: .exe mainPicture logoPicture outputPicture");
        return -1;
    }

    AVFrame *main_frame = get_frame_from_jpegfile(argv[1]);
    AVFrame *logo_frame = get_frame_from_jpegfile(argv[2]);
    AVFrame *result_frame = av_frame_alloc();
    int ret = 0;

    if ((ret = init_filters(main_frame, logo_frame, 50, 50)) < 0) {
        printf("Init filters failed.\n");
        goto FAILED;
    }

    if (main_picture_mix_logo(main_frame, logo_frame, result_frame) < 0) {
        printf("main_picture_mix_logo failed.\n");
        goto FAILED;
    }

    save_frame_to_jpeg(argv[3], result_frame);

FAILED:
    if (main_frame != NULL)
        av_frame_free(&main_frame);

    if (logo_frame != NULL)
        av_frame_free(&logo_frame);

    if (result_frame != NULL)
        av_frame_free(&result_frame);

    if (filter_graph != NULL)
        avfilter_graph_free(&filter_graph);
    printf("Finish.\n");

    return 0;
}
